import numpy as np
import regression

def drawPlot(xArr, yArr, yHat, k, pngName):
    xMat = np.matrix(xArr)
    strInd = xMat[:, 1].argsort(0)
    xSort = xMat[strInd][:, 0, :]
    import matplotlib.pyplot as plt
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.plot(xSort[:, 1], yHat[strInd])
    ax.scatter(xMat[:, 1].flatten().A[0], np.matrix(yArr).T.flatten().A[0], s=2, c='red')
    plt.title("LWLR k=%f" %(k))
    plt.savefig(pngName)

def lwlrTestAndDraw(xArr, yArr, k):
    yHat = regression.lwlrTest(xArr, xArr, yArr, k)
    drawPlot(xArr, yArr, yHat, k, 'lwlr_k%f.png' %(k))

# 分别使用不同的k值进行lwlr回归，并绘制拟合曲线和原始值散点图，观察欠拟合和过拟合

def drawTests():
    xArr, yArr = regression.loadDataSet('ex0.txt')

    lwlrTestAndDraw(xArr, yArr, 1.0)
    lwlrTestAndDraw(xArr, yArr, 0.05)
    lwlrTestAndDraw(xArr, yArr, 0.01)    
    lwlrTestAndDraw(xArr, yArr, 0.003)        

